import pandas as pd
import os

# --- 1. 定义文件名 ---
input_csv_file = "ocr_results.csv"
output_csv_file = "ocr_results_merged.csv"

# --- 2. 检查输入文件是否存在 ---
if not os.path.exists(input_csv_file):
    print(f"错误: 输入文件未找到: {input_csv_file}")
    exit()

print(f"正在读取 {input_csv_file}...")
try:
    df = pd.read_csv(input_csv_file)
except pd.errors.EmptyDataError:
    print(f"文件 {input_csv_file} 为空，无法处理。")
    exit()

# --- 3. 分离成功的数据和错误/状态数据 ---

# 成功的数据 (text_rank 为 1 或 2)
df_success = df[df["text_rank"] > 0].copy()

# 错误/状态数据 (text_rank 为 0, 例如 "No label image found")
df_status = df[df["text_rank"] == 0].copy()


# --- 4. 处理成功的数据：执行 'pivot' 操作 ---
df_pivoted = pd.DataFrame()
if not df_success.empty:
    # 将 'text_rank' (1, 2) 转换为列名
    # index='file_path':       让 'file_path' 保持为行索引
    # columns='text_rank':     将这一列的值 (1, 2) 转换为新的列
    # values=[...]:            这些是要在新列中显示的值
    df_pivoted = df_success.pivot(
        index="file_path",
        columns="text_rank",
        values=["text", "confidence", "y_coordinate", "full_bbox"],
    )

    # 'pivot' 会创建多级列名 (例如 ('text', 1), ('confidence', 1))
    # 我们需要将它们“扁平化”为单级列名 (例如 'text_1', 'confidence_1')

    # 修改：移除了 inplace=True 并重新赋值
    # 修改：将 levels 和 level 都指定为列表
    df_pivoted.columns = df_pivoted.columns.set_levels(
        [df_pivoted.columns.levels[1].astype(int)], level=[1]
    )

    # 扁平化
    df_pivoted.columns = [f"{col[0]}_{int(col[1])}" for col in df_pivoted.columns]

    # 将 'file_path' 从索引变回普通列
    df_pivoted = df_pivoted.reset_index()

    # 添加一个 'status' 列
    df_pivoted["status"] = "Success"
else:
    print("未找到任何成功的 OCR 结果 (rank 1 or 2)。")

# --- 5. 处理错误/状态数据 ---
df_status_formatted = pd.DataFrame()
if not df_status.empty:
    # 我们只需要 file_path 和 text (其中包含错误信息)
    df_status_formatted = df_status[["file_path", "text"]].copy()

    # 重命名 'text' 列为 'status' 或 'error_message'
    df_status_formatted = df_status_formatted.rename(columns={"text": "status"})
else:
    print("未找到任何状态/错误行 (rank 0)。")

# --- 6. 合并成功的数据和错误数据 ---

# pd.concat 会将两个 DataFrame 堆叠在一起
# 'sort=False' 保持列的原始顺序
df_final = pd.concat([df_pivoted, df_status_formatted], ignore_index=True, sort=False)

# 填充 'NaN' 值 (例如，错误行没有 'text_1', 成功行没有 'error_message')
# 我们用空字符串 '' 填充
df_final = df_final.fillna("")

# --- 7. 重新排序，使列的顺序更合理 ---
expected_cols_base = [
    "file_path",
    "status",
    "text_1",
    "confidence_1",
    "y_coordinate_1",
    "full_bbox_1",
    "text_2",
    "confidence_2",
    "y_coordinate_2",
    "full_bbox_2",
]

# 确保所有预期的列都存在（即使在原始数据中没有 rank 2 的情况）
for col in expected_cols_base:
    if col not in df_final.columns:
        df_final[col] = ""

# --- 7.5 (新增) 创建 'ocr_name' 列 ---
# 在这里创建，可以确保 'text_1' 和 'text_2' 都存在并且是字符串 (不是 NaN)
print("正在创建 'ocr_name' 列...")
df_final["ocr_name"] = df_final["text_1"] + '_'+df_final["text_2"]

# --- (修改) 定义最终的列顺序，包含 'ocr_name' ---
final_ordered_cols = [
    "file_path",
    "status",
    "ocr_name",  # <-- 新增的列
    "text_1",
    "confidence_1",
    "y_coordinate_1",
    "full_bbox_1",
    "text_2",
    "confidence_2",
    "y_coordinate_2",
    "full_bbox_2",
]


# 按预期顺序排列
df_final = df_final[final_ordered_cols]

# --- 8. 保存到新的 CSV 文件 ---
df_final.to_csv(output_csv_file, index=False, encoding="utf-8")

print("\n=============================================")
print(f"处理完成！合并后的文件已保存到: {output_csv_file}")
print(f"总共处理了 {len(df_final)} 个独立的文件路径。")
